# @Time : 2025-01-17 10:22
# @Author : Fioman 
# @Phone : 13149920693
"""
寻找物料区域
"""
import os
import time

import cv2

from common.exception_helper import error_handler_vision
from vision_process.calc_config import *
from vision_process.vision_helper import *


def contour_is_ok(cnt, w):
    """
    轮廓是否满足要求,除了轮廓的面积之外,轮廓也不能太靠近边缘,如果轮廓的中心离边缘太近,
    也被视为无效轮廓.
    :param cnt:
    :param w:
    :return:
    """
    if cv.contourArea(cnt) <= filterArea:
        return False
    contourCenter = cv.minAreaRect(cnt)[0]
    if contourCenter[0] < int(50 * cp.mm2pix) or contourCenter[0] > (w - int(50 * cp.mm2pix)):
        return False
    return True


@error_handler_vision
def find_totalboard(image, fileIndex=None) -> calc_config.TotalboardResult:
    result = TotalboardResult()
    thresUsed = image[:150, :]
    thresUsedVal = thresUsedTotalboard
    image_show(f"TotalboardThresUsed_({thresUsedVal})", thresUsed)
    # 1. 先做阈值分割,初步查找物料区域
    _, board = cv.threshold(image, thresUsedVal, 255, cv.THRESH_BINARY)
    image_show(f"BoardThres_{thresUsedVal}", board)
    openKernelSize = (3, 3)
    openKernel = cv.getStructuringElement(cv.MORPH_RECT, openKernelSize)
    board = cv.morphologyEx(board, cv.MORPH_OPEN, openKernel, iterations=3)
    image_show(f"BoardOpened_(3,3)_3", board)

    cnts, _ = cv.findContours(board.copy(), cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
    cnts = [cnt for cnt in cnts if cv.contourArea(cnt) > calc_config.filterArea]
    print(f"轮廓个数: {len(cnts)}")
    if len(cnts) == 0:
        result.info = "找到的轮廓的个数为0,物料区域查找失败"
        return result

    # 2. 将轮廓填充为白色
    # 创建掩码图像
    mask = np.zeros_like(board)
    # 填充所有轮廓
    cv.fillPoly(mask, cnts, colorWhite)
    # 将填充后的掩码应用到原图
    board = cv.bitwise_and(board, mask)
    image_show("ContourFilled", board)

    # 3. 闭操作扩大轮廓
    closeKernelSize = (5, 5)
    closeIterations = 8
    closeKernel = cv.getStructuringElement(cv.MORPH_RECT, closeKernelSize)
    board = cv.morphologyEx(board, cv2.MORPH_CLOSE, closeKernel, iterations=closeIterations)
    image_show(f"BoardClosed_{closeIterations}", board)

    # 4. 做完闭操作之后,再找一次轮廓
    cnts, _ = cv.findContours(board.copy(), cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
    validContours = [cnt for cnt in cnts if contour_is_ok(cnt, image.shape[1])]
    if len(validContours) == 0:
        result.info = "轮廓的个数为0,物料区域查找失败"
    if len(validContours) >= 2:
        for index, contour in enumerate(validContours):
            anotherContour = validContours[(index + 1) % len(validContours)]
            point1 = int(contour[0][0][0]), int(contour[0][0][1])
            point2 = int(anotherContour[0][0][0]), int(anotherContour[0][0][1])
            cv.line(board, point1, point2, colorWhite, 1)
            image_show(f"JointContour", board)
        cnts, _ = cv.findContours(board.copy(), cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
        cntFirstFinded = sorted(cnts, key=cv.contourArea, reverse=True)[0]
    else:
        cntFirstFinded = validContours[0]

    # 5. 找到轮廓的四个顶点
    boxMaster = cv.minAreaRect(cntFirstFinded)
    rectPoints = np.intp(cv.boxPoints(boxMaster))
    leftTop, rightTop, rightBottom, leftBottom = get_rect_points_clockwise(rectPoints)
    # 6. 获取轮廓和矩形的近似度
    rectRatio = get_contour_is_rect_ratio(cntFirstFinded)
    print(f"第一次找到的轮廓和最小外接矩形的近似度为: {rectRatio}")
    if rectRatio < 0.97:
        # 如果小于0.97 就再找一次
        cv.destroyAllWindows()
        moveWidth = int(get_two_point_distance(leftTop, rightTop) * 0.15)
        moveHeight = int(get_two_point_distance(leftTop, leftBottom) * 0.15)
        extend = 30
        boardFinal = board.copy()
        x, y = leftTop
        boardFinal[y - extend:y + moveHeight, x - extend:x + moveWidth] = 0
        x, y = rightTop
        boardFinal[y - extend:y + moveHeight, x - moveWidth:x + extend] = 0
        x, y = leftBottom
        boardFinal[y - moveHeight:y + extend, x - extend:x + moveWidth] = 0
        x, y = rightBottom
        boardFinal[y - moveHeight:y + extend, x - moveWidth:x + extend] = 0
        image_show("FinalBoardAngleMoved", boardFinal)
        cnts, _ = cv.findContours(boardFinal.copy(), cv.RETR_EXTERNAL,
                                  cv.CHAIN_APPROX_SIMPLE)
        cntFinalFinded = sorted(cnts, key=cv.contourArea, reverse=True)[0]
        boxMaster = cv.minAreaRect(cntFirstFinded)
        rectPoints = np.intp(cv.boxPoints(boxMaster))
    else:
        cntFinalFinded = cntFirstFinded
    totalboardThres = cv.drawContours(mask.copy(), cntFinalFinded, -1, colorWhite, -1)
    image_show("TotalboardThres", totalboardThres)

    # 然后就是检测box的角度问题,宽和高跟角度有关系
    (cx, cy), (w, h), angle = boxMaster
    print(f"boxMaster角度: {angle}")
    if abs(angle) > 45:
        w, h = h, w
        angle = angle + 90 if angle < 0 else angle - 90
        boxMaster = ((cx, cy), (w, h), angle)

    imageColor = cv.cvtColor(image, cv.COLOR_GRAY2RGB)
    if fileIndex is not None:
        writePos = int(cx), int(cy)
        cv.putText(imageColor, str(fileIndex), writePos, cv.FONT_HERSHEY_SIMPLEX,
                   5, colorBlue, 5)
    cv.drawContours(imageColor, [rectPoints], -1, colorGreen, 4)
    image_show("TotalboardRes", imageColor)
    cv.destroyAllWindows()
    totalboard, offset = get_board_without_angle(image, boxMaster)
    totalboardRes, _ = get_board_without_angle(imageColor, boxMaster, extend=30)
    result.totalboardThres = totalboardThres
    result.totalboardRes = totalboardRes
    result.totalboardGray = totalboard
    result.totalboardColor = totalboardRes
    result.totalboardOffset = offset
    result.boxMaster = boxMaster
    result.rectPoints = rectPoints
    result.state = True
    return result


def put_calc_info_to_image(image, rectPoints, calcSize):
    """
    :param image:
    :param rectPoints:
    :param calcSize:
    :return:
    """

    cv.drawContours(image, [rectPoints],
                    -1, colorGreen, 3)
    writePos = (int(image.shape[1] / 2), int(image.shape[0] / 2))
    sizeInfo = f"W:{calcSize[0]}, H:{calcSize[1]}"
    cv.putText(image, sizeInfo, (int(writePos[0] - 200), int(writePos[1] + 200)),
               cv.FONT_HERSHEY_SIMPLEX, 2, colorRed, 3)
    # imageRes, _ = get_board_without_angle(image, box, extend=100)
    image_show("TotalboardWithInfo", image)
    return image


def check_totalboard_is_ok(image) -> TotalboardResult:
    """
    检测物料区域是否OK,来料是否正常
    :param image:
    :return:
    """
    result: TotalboardResult = find_totalboard(image)
    if not result.state:
        return result

    # 计算物料区域到传感器的距离
    w, h = result.boxMaster[1]
    widthChecked, heightChecked = round(w / cp.mm2pix, 2), round(h / cp.mm2pix, 2)
    imageColor = cv.cvtColor(image.copy(), cv.COLOR_GRAY2BGR)
    totalboardRes = put_calc_info_to_image(imageColor, result.rectPoints,
                                           (widthChecked, heightChecked))
    result.totalboardColor = totalboardRes
    heighMin = cp.heightMin

    if not   heighMin <= heightChecked <= cp.heightMax:
        result.state = False
        result.info = (f"物料检测未通过,物料高度超出有效范围,检测高度:{heightChecked},有效范围:"
                       f"({heighMin},{cp.heightMax})")
        return result
    result.state = True
    result.info = (f"物料检测通过,物料宽:{widthChecked},物料高:{heightChecked},"
                   f"高度ok范围:({cp.heightMin},{cp.heightMax})")
    return result


if __name__ == '__main__':
    filePath = r"F:\raw_2025\zhanhua"
    keepDir = filePath
    if not os.path.isdir(keepDir):
        keepDir, _ = os.path.split(filePath)

    keepOkPath = os.path.join(keepDir, "ok")
    keepOkResPath = os.path.join(keepDir, "ok_res")
    keepFailPath = os.path.join(keepDir, "fail")
    keepFailResPath = os.path.join(keepDir, "fail_res")

    dirList = [keepOkPath, keepOkResPath, keepFailPath, keepFailResPath]
    for dirDemo in dirList:
        if not os.path.exists(dirDemo):
            os.makedirs(dirDemo)

    fileNames = []
    for root, dirs, files in os.walk(keepDir):
        if root != keepDir:
            continue
        for file in files:
            if not file.endswith(".bmp"):
                continue
            fileNames.append(file)

    for indexFile, fileName in enumerate(fileNames, start=1):
        filePathReal = os.path.join(keepDir, fileName)
        imageSrc = cv.imread(filePathReal, cv.IMREAD_GRAYSCALE)
        if imageSrc is None:
            print(f"{'*' * 10} 第 {indexFile} 张图 ({fileName}),图像读取错误,已跳过. {'*' * 10}")
            continue
        else:
            print(f"{'*' * 10} 第 {indexFile} 张图 ({fileName}),开始识别 {'*' * 10}")
        calcStart = time.time()
        resultTest: calc_config.TotalboardResult = check_totalboard_is_ok(imageSrc)
        if resultTest.state:
            print(f"物料区域查找成功,算法耗时: {(time.time() - calcStart):.4f}")
            os.remove(filePathReal)
            cv.imwrite(os.path.join(keepOkPath, fileName), imageSrc)
            imageRes = get_size_scale_image(resultTest.totalboardColor, sizeScale=3)
            cv.imwrite(os.path.join(keepOkResPath, fileName), imageRes)
        else:
            print(f"物料区域查找失败XXXX,原因: {resultTest.info},算法耗时: {(time.time() - calcStart):.4f}")
            cv.imwrite(os.path.join(keepFailPath, fileName), imageSrc)
            imageRes = get_size_scale_image(resultTest.totalboardColor, sizeScale=3)
            cv.imwrite(os.path.join(keepFailResPath, fileName), imageRes)
            os.remove(filePathReal)
